# -*- coding: utf-8 -*-
"""
Created on Fri Aug 21 11:36:39 2020

@author: Farman
"""

import pygrib
import numpy
import joblib
import sys
import pathlib


def update_ma_data(masked_array):
    '''
    Set masked values to zero for further
    '''
    m = masked_array.mask
    d = masked_array.data * (m ^ True)
    masked_array = numpy.ma.masked_array(data=d, mask=m)
    return masked_array


def show_messages(grib_file):
    gribs = pygrib.open(grib_file)
    messages = gribs.read()
    gribs.close()
    
    print('\nMessages in file : %s'%grib_file)
    print('-' * 100)
    
    for m in messages:
        print(m)
    
    print()
    return


def get_grib_data(grib_file, layers):
    gribs = pygrib.open(grib_file)
    messages = gribs.read()
    #print(len(messages))
    gribs.close()
    
    if isinstance(layers, int):
        data = messages[layers].data()
    else:
        data = []
        
        for layer in layers:
            d = messages[layer].data()
            data.append(d)
    
    ma_list = []
    
    for d in data:
        if isinstance(d, numpy.ma.masked_array):
            ma_list.append(update_ma_data(d))
        else:
            ma_list.append(d)
        
    return ma_list


# get_era5_t_e_vsm() and get_era5_t_e_vsm2()
# are two different method of data processing.
# on PC, both of them almost are same in performance.
# use one of them for your convinience.

def get_era5_t_e_vsm(era5_grib_file):
    '''
    Get temperature, evaporation and volumetric soil water data from file of
    path 'ERA5'.
    
    Temperature in unit of K, mean value of 24 hours of day.
    Evaporation in unit of m, accumulated value of 24 hours of day.
    Volumetric soil water in m^3 / m^3, upper layer and deeper layer.
    the upper layer is the mean value of original layer 1 and 2 of 24 hours of day,
    the deeper layer is the mean value of original layer 3 and 4 of 24 hours of day.
    '''
    gribs = pygrib.open(era5_grib_file)
      
    print('Reading data of hour %2d ...'%0)
    
    df = gribs.read(1)[0].data()[0]
    k = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    e = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    u = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    u += df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    d = df.data * (True ^ df.mask)
    
    df = gribs.read(1)[0].data()[0]
    d = df.data * (True ^ df.mask)
    
    for hour in range(1, gribs.messages // 6):
        print('Reading data of hour %2d ...'%hour)
        
        df = gribs.read(1)[0].data()[0]
        k += df.data * (True ^ df.mask)
    
        df = gribs.read(1)[0].data()[0]
        e += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        u += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        u += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        d += df.data * (True ^ df.mask)
        
        df = gribs.read(1)[0].data()[0]
        d += df.data * (True ^ df.mask)
        
    gribs.close()
        
    k = k / 24
    u = u / 48
    d = d / 48
    return k, e, u, d


def get_era5_t_e_vsm2(era5_grib_file):
    '''
    Get temperature, evaporation and volumetric soil water data from file of
    path 'ERA5'.
    
    Temperature in unit of K, mean value of 24 hours of day.
    Evaporation in unit of m, accumulated value of 24 hours of day.
    Volumetric soil water in m^3 / m^3, upper layer and deeper layer.
    the upper layer is the mean value of original layer 1 and 2 of 24 hours of day,
    the deeper layer is the mean value of original layer 3 and 4 of 24 hours of day.
    '''
    gribs = pygrib.open(era5_grib_file)
    messages = gribs.read()
    gribs.close()
    
    print('Reading data of hour %2d ...'%0)
    
    df = messages[0].data()[0]
    k = df.data * (True ^ df.mask)
    
    df = messages[1].data()[0]
    e = df.data * (True ^ df.mask)
    
    df = messages[2].data()[0]
    u = df.data * (True ^ df.mask)
    
    df = messages[3].data()[0]
    u += df.data * (True ^ df.mask)
    
    df = messages[4].data()[0]
    d = df.data * (True ^ df.mask)
    
    df = messages[5].data()[0]
    d = df.data * (True ^ df.mask)
    
    for hour in range(1, gribs.messages // 6):
        print('Reading data of hour %2d ...'%hour)
        
        df = messages[hour * 6 + 0].data()[0]
        k += df.data * (True ^ df.mask)
    
        df = messages[hour * 6 + 1].data()[0]
        e += df.data * (True ^ df.mask)
        
        df = messages[hour * 6 + 2].data()[0]
        u += df.data * (True ^ df.mask)
        
        df = messages[hour * 6 + 3].data()[0]
        u += df.data * (True ^ df.mask)
        
        df = messages[hour * 6 + 4].data()[0]
        d += df.data * (True ^ df.mask)
        
        df = messages[hour * 6 + 5].data()[0]
        d += df.data * (True ^ df.mask)
        
    k = k / 24
    u = u / 48
    d = d / 48
    return k, e, u, d


def tran_day_data_of_era5_t_e_vsm(era5_grib_file, dest_dump_file):
    if pathlib.Path(dest_dump_file).exists():
        print('Dest file exist, skip : %s'%dest_dump_file)
        return
    
    print('\nReading from file : %s'%era5_grib_file)
    print('-' * 40)
    k, e, u, d = get_era5_t_e_vsm(era5_grib_file)
    print('Dumping to   file : %s'%dest_dump_file)
    joblib.dump((k, e, u, d), dest_dump_file)
    print('-' * 10, 'Done', '-' * 10, '\n')
    return


#-----------------------------------------------------------------------------W


if __name__ == '__main__':
    if len(sys.argv) != 3:
        print('\nUsage')
        print('-' * 40)
        print('python GribVsmTran.py era5_grib_file_of_day dest_dump_file')
        print()
    else:
        era5_grib_file, dest_dump_file = sys.argv[1:]
        tran_day_data_of_era5_t_e_vsm(era5_grib_file, dest_dump_file)

    
'''
era5_grib_file = '/media/farman/Data/CUIT-NewFarman/Data/ERA5/2018_06_28.grib'
dump_file = '1.joblib'

tran_day_data_of_era5_t_e_vsm(era5_grib_file, dump_file)

import time

start = time.perf_counter()
k, e, u, d = get_era5_t_e_vsm(era5_grib_file)
print('Time used %s'%(time.perf_counter() - start))

start = time.perf_counter()
k, e, u, d = get_era5_t_e_vsm2(era5_grib_file)
print('Time used %s'%(time.perf_counter() - start))


#d = get_grib_data(era5_grib_file, 6)
#d = get_era5_t_e_vsm(grib_file)
#'''
